from call_llm import _run_all_configs_once
import asyncio
from typing import List, Optional
import re


async def rewrite_doubt_by_llm(
    doubt_list: List[str],  # 接收字符串列表
    concurrency_limit: int = 5,
    name: Optional[str] = None,
    system_prompt: Optional[str] = None,
    user_prompt: Optional[str] = None,  # 这个暂时用不上的
    concurrent: Optional[int] = None,
) -> List[str]:
    """并发改写疑问句：
    - 若文本包含“规定/条例/条文”，直接原样返回
    - 否则调用大模型进行改写
    - 并发量由 concurrency_limit 控制；模型参数可通过可选参数传入
    """

    if not system_prompt:
        system_prompt = """
        您是一个专业的纪律审查改写专家，专门从违纪案例描述中提取核心要素，改写能够精准匹配处分条款库的检索查询。
        ## 要求
        - 准确提取描述中的核心行为
        - 提取行为造成的主要结果或影响
        - 使用通用、规范化的法律或纪律术语，而非案件叙述性语言
        - 不包含具体人名、单位名、金额或项目细节
        - 然后输出3个不同query改写结果（一共30字内），然后合并在一起，之前用分号“;”隔开。
        ## 思考模式
        /no_think
        """

    # 简单关键词判断函数
    def is_regulation(text: str) -> bool:
        if not isinstance(text, str):
            return False
        return any(k in text for k in ("规定", "条例", "条文", "_"))

    semaphore = asyncio.Semaphore(max(1, int(concurrency_limit)))

    # async def process_item(text: str, system_prompt, user_prompt) -> str:
    async def process_item(text: str, system_prompt, user_prompt, delay: float = 0) -> str:
        if not user_prompt:
            user_prompt = f"{text}"
        # 命中关键词：直接返回
        if is_regulation(text):
            print(f"关键词匹配，直接返回: {text}")
            return text
        # 未命中：调用大模型
        async with semaphore:
            # 添加延迟避免并发请求过于密集
            if delay > 0:
                await asyncio.sleep(delay)
            print(f"开始调用LLM改写: {text}")
            # 默认将当前文本作为 user_prompt；若外部传入 user_prompt，则按传入值使用
            results = await _run_all_configs_once(
                name=name,
                system_prompt=system_prompt or "you are a helpful assistant.",
                user_prompt=(user_prompt if user_prompt is not None else text),
                concurrent=concurrent or 1,
            )
            # _run_all_configs_once 返回 List[str]，取第一条作为该项结果
            result = results[0] if results else text  # 报错返回None，返回原文本
            # 去掉<think>\n\n</think>\n\n
            result = re.sub(r'<think>.*?</think>', '', result, flags=re.DOTALL).strip()
            print(f"LLM改写结果: {result}")
            return result

    # 每个请求递增延迟，避免并发过于密集导致超时 (可选: 0.01秒/0.05秒/0.1秒/0.2秒)
    tasks = [process_item(t, system_prompt, user_prompt, delay=i*0.01) for i, t in enumerate(doubt_list)]
    # tasks = [process_item(t, system_prompt, user_prompt) for t in doubt_list]  # 原始版本，无延迟
    # tasks = [process_item(t, system_prompt, user_prompt, delay=i*0.05) for i, t in enumerate(doubt_list)]
    return await asyncio.gather(*tasks)


if __name__ == "__main__":
    # 示例：包含命中关键词与非命中两类
    sample = [
    "张务实同志违反政治纪律，理想信仰动摇，搞封建迷信活动。",
    "_",
    "张务实同志违反廉洁纪律，以权谋私，收受可能影响公正执行公务的礼金、礼品。",
    "_",
    "张务实同志违反工作纪律，履行管党治党责任不力，工作中存在形式主义、官僚主义问题。",
    "_",
    "张务实同志涉嫌滥用职权，为不法商人‘开绿灯’，造成国家资产重大损失。",
    "_"
    ]
    joined_sample = "|".join(sample)
    out = asyncio.run(rewrite_doubt_by_llm(sample, concurrency_limit=3))
    print(out)
